/*
 * Copyright (c) 2025 Huawei Technologies Co., Ltd.
 * This file is a part of the CANN Open Software.
 * Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
 * Please refer to the License for details. You may not use this file except in compliance with the License.
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
 * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
 * See LICENSE in the root of the software repository for the full text of the License.
 */

#include "../../common/data_utils.h"
#include "./op_host/adjust_softmax_res_custom_tiling.h"
#ifndef ASCENDC_CPU_DEBUG
#include "acl/acl.h"
#include "aclrtlaunch_adjust_softmax_res_custom.h"
#include "tiling/platform/platform_ascendc.h"
#else
#include "tikicpulib.h"
extern "C" __global__ __aicore__ void adjust_softmax_res_custom(GM_ADDR srcGm, GM_ADDR dstGm, GM_ADDR max, 
                                                                AdjustSoftMaxResTilingData tiling);
#endif

namespace {
constexpr uint32_t USED_CORE_NUM = 1;
constexpr uint32_t TILING_SIZE = 4;
constexpr uint32_t SRC_HEIGHT = 32;
constexpr uint32_t SRC_WIDTH = 32;
}

extern void GenerateTiling(uint32_t height, uint32_t width, uint8_t* tilingBuf);

static bool CompareResult(const void *outputData, int64_t outSize) {
    void *goldenData;
#ifdef ASCENDC_CPU_DEBUG
    goldenData = (uint8_t *)AscendC::GmAlloc(outSize);
#else
    CHECK_ACL(aclrtMallocHost((void **)(&goldenData), outSize));
#endif
    size_t goldenSize = outSize;
    bool ret = ReadFile("../output/golden.bin", goldenSize, goldenData, goldenSize);
    if (ret) {
        printf("ReadFile golden.bin success!\n");
    } else {
        printf("test failed!\n");
        return false;
    }
    constexpr float EPS = 1e-4;
    int64_t wrongNum = 0;

    for (int i = 0; i < outSize / sizeof(float); i++) {
        float a = (reinterpret_cast<const float *>(outputData))[i];
        float b = (reinterpret_cast<const float *>(goldenData))[i];
        float ae = std::abs(a - b);
        float re = ae / abs(b);
        if (ae > EPS && re > EPS) {
            wrongNum++;
        }
    }
#ifdef ASCENDC_CPU_DEBUG
    AscendC::GmFree((void *)goldenData);
#else
    CHECK_ACL(aclrtFreeHost(goldenData));
#endif
    if (wrongNum != 0) {
        printf("wrongNum: %ld\n", wrongNum);
        return false;
    } else {
        printf("CompareResult golden.bin success!\n");
        return true;
    }
}

int32_t main(int32_t argc, char *argv[]) {
    size_t tilingSize = TILING_SIZE * sizeof(uint32_t);
    size_t inputSize = SRC_HEIGHT * SRC_WIDTH * sizeof(uint32_t);
    size_t inputMaxSize = SRC_HEIGHT * 8 * sizeof(uint32_t);
    size_t outputSize = SRC_HEIGHT * SRC_WIDTH * sizeof(uint32_t);
#ifdef ASCENDC_CPU_DEBUG
    uint8_t *x = (uint8_t *)AscendC::GmAlloc(inputSize);
    uint8_t *y = (uint8_t *)AscendC::GmAlloc(outputSize);
    uint8_t *max = (uint8_t *)AscendC::GmAlloc(inputMaxSize);
    uint8_t *tiling = (uint8_t *)AscendC::GmAlloc(tilingSize);

    ReadFile("../input/input_softmax.bin", inputSize, x, inputSize);
    ReadFile("../input/input_max.bin", inputMaxSize, max, inputMaxSize);

    GenerateTiling(SRC_HEIGHT, SRC_WIDTH, tiling);

    AscendC::SetKernelMode(KernelMode::AIV_MODE); // run in aiv mode

    ICPU_RUN_KF(adjust_softmax_res_custom, 1, x, y, max, 
                *reinterpret_cast<AdjustSoftMaxResTilingData *>(tiling)); // use this macro for cpu debug

    WriteFile("../output/output.bin", y, outputSize);

    bool goldenResult = true;
    goldenResult = CompareResult(y, outputSize);

    AscendC::GmFree((void *)x);
    AscendC::GmFree((void *)y);
    AscendC::GmFree((void *)tiling);
#else
    CHECK_ACL(aclInit(nullptr));
    int32_t deviceId = 0;
    CHECK_ACL(aclrtSetDevice(deviceId));
    aclrtStream stream = nullptr;
    CHECK_ACL(aclrtCreateStream(&stream));

    uint8_t *xHost; 
    uint8_t *yHost; 
    uint8_t *maxHost; 
    uint8_t *tiling;
    uint8_t *xDevice; 
    uint8_t *yDevice; 
    uint8_t *maxDevice;

    CHECK_ACL(aclrtMallocHost((void **)(&xHost), inputSize));
    CHECK_ACL(aclrtMallocHost((void **)(&yHost), outputSize));
    CHECK_ACL(aclrtMallocHost((void **)(&maxHost), inputMaxSize));
    CHECK_ACL(aclrtMallocHost((void **)(&tiling), tilingSize))

    CHECK_ACL(aclrtMalloc((void **)&xDevice, inputSize, ACL_MEM_MALLOC_HUGE_FIRST));
    CHECK_ACL(aclrtMalloc((void **)&yDevice, outputSize, ACL_MEM_MALLOC_HUGE_FIRST));
    CHECK_ACL(aclrtMalloc((void **)&maxDevice, inputMaxSize, ACL_MEM_MALLOC_HUGE_FIRST));

    ReadFile("../input/input_softmax.bin", inputSize, xHost, inputSize);
    ReadFile("../input/input_max.bin", inputMaxSize, maxHost, inputMaxSize);

    GenerateTiling(SRC_HEIGHT, SRC_WIDTH, tiling);

    // Copy host memory to device memory
    CHECK_ACL(aclrtMemcpy(xDevice, inputSize, xHost, inputSize, ACL_MEMCPY_HOST_TO_DEVICE));
    CHECK_ACL(aclrtMemcpy(maxDevice, inputMaxSize, maxHost, inputMaxSize, ACL_MEMCPY_HOST_TO_DEVICE));

    // Execute the kernel
    ACLRT_LAUNCH_KERNEL(adjust_softmax_res_custom)
    (USED_CORE_NUM, stream, xDevice, yDevice, maxDevice, reinterpret_cast<AdjustSoftMaxResTilingData *>(tiling));

    // Wait for the stop event to complete
    CHECK_ACL(aclrtSynchronizeStream(stream));

    // Copy result to host memory and write to output file
    CHECK_ACL(aclrtMemcpy(yHost, outputSize, yDevice, outputSize, ACL_MEMCPY_DEVICE_TO_HOST));
    WriteFile("../output/output.bin", yHost, outputSize);

    // Compare the result with the golden result
    bool goldenResult = true;
    goldenResult = CompareResult(yHost, outputSize);

    // Clean up memory
    CHECK_ACL(aclrtFree(xDevice));
    CHECK_ACL(aclrtFree(yDevice));
    CHECK_ACL(aclrtFree(maxDevice));
    CHECK_ACL(aclrtFreeHost(xHost));
    CHECK_ACL(aclrtFreeHost(yHost));
    CHECK_ACL(aclrtFreeHost(maxHost));
    CHECK_ACL(aclrtFreeHost(tiling));

    CHECK_ACL(aclrtDestroyStream(stream));
    CHECK_ACL(aclrtResetDevice(deviceId));
    CHECK_ACL(aclFinalize());
#endif

    if (goldenResult) {
        printf("test pass!\n");
    } else {
        printf("test failed!\n");
    }
    return 0;
}